import torch
import torch.nn as nn

class TransformerModel(nn.Module):
    def __init__(self, emb_dim=204, n_heads=6, hid_dim=512, n_layers=6, output_dim=204, dropout=0.1):
        super().__init__()

        self.pos_encoding = PositionalEncoding(emb_dim, dropout)
        self.transformer_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(emb_dim, n_heads, hid_dim, dropout), n_layers
        )
        self.fc = nn.Sequential(nn.Linear(emb_dim, 32),
                                nn.Linear(32, output_dim))
        self.dropout = nn.Dropout(dropout)
        self.softmax = nn.Softmax(dim=-1)
 
    def forward(self, src):
        src = torch.transpose(src, 0, 1)
        embedded = self.pos_encoding(src)
        outputs = self.transformer_encoder(embedded)
        last_output = outputs[-1, :, :]
        last_output = self.dropout(last_output)
        last_output = self.fc(last_output)
        last_output = last_output.view((last_output.shape[0], 6, 34))
        last_output = self.softmax(last_output)
        return last_output

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=50):
        
        super(PositionalEncoding, self).__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term2 = torch.pow(torch.tensor(10000.0),torch.arange(0, d_model, 2).float()/d_model)
        div_term1 = torch.pow(torch.tensor(10000.0),torch.arange(1, d_model, 2).float()/d_model)
        
        pe[:, 0::2] = torch.sin(position * div_term2)
        pe[:, 1::2] = torch.cos(position * div_term1)
        
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)
    
if __name__ == "__main__":
    model = TransformerModel()
    inputs = torch.rand(1, 50, 204)
    outputs = model(inputs)
    print(outputs)